#include <rclcpp/rclcpp.hpp>
#include <std_msgs/msg/string.hpp>
#include <opentelemetry/trace/provider.h>
#include <opentelemetry/exporters/ostream/span_exporter.h>
#include <opentelemetry/sdk/trace/simple_processor.h>
#include <opentelemetry/sdk/trace/tracer_provider.h>
#include <opentelemetry/sdk/trace/samplers/always_on.h>
#include <opentelemetry/sdk/resource/resource.h>
#include <opentelemetry/nostd/shared_ptr.h>
#include <opentelemetry/nostd/unique_ptr.h>
#include <opentelemetry/trace/context.h>
#include <opentelemetry/trace/propagation/http_trace_context.h>
#include <opentelemetry/trace/span_context.h>
#include <opentelemetry/trace/span_id.h>
#include <opentelemetry/trace/trace_id.h>

#include <chrono>
#include <thread>
#include <string>

namespace trace_api = opentelemetry::trace;
namespace trace_sdk = opentelemetry::sdk::trace;
namespace resource_sdk = opentelemetry::sdk::resource;
namespace nostd = opentelemetry::nostd;

// Initialize OpenTelemetry tracing
void InitTracer() {
    auto exporter = std::unique_ptr<trace_sdk::SpanExporter>(
        new opentelemetry::exporter::trace::OStreamSpanExporter);
    
    auto processor = std::unique_ptr<trace_sdk::SpanProcessor>(
        new trace_sdk::SimpleSpanProcessor(std::move(exporter)));
    
    auto resource = resource_sdk::Resource::Create({
        {"service.name", "ros2-node-c"},
        {"service.version", "1.0.0"},
        {"service.instance.id", "node-c-001"}
    });
    
    auto sampler = std::unique_ptr<trace_sdk::Sampler>(
        new trace_sdk::AlwaysOnSampler);
    
    auto provider = nostd::shared_ptr<trace_api::TracerProvider>(
        new trace_sdk::TracerProvider(std::move(processor), resource, std::move(sampler)));
    
    trace_api::Provider::SetTracerProvider(provider);
}

class NodeC : public rclcpp::Node {
public:
    NodeC() : Node("node_c") {
        // Initialize OpenTelemetry
        InitTracer();
        tracer_ = trace_api::Provider::GetTracerProvider()->GetTracer("node-c");
        
        // Create subscriber
        subscription_ = this->create_subscription<std_msgs::msg::String>(
            "topic_b_to_c", 10,
            std::bind(&NodeC::messageCallback, this, std::placeholders::_1));
        
        RCLCPP_INFO(this->get_logger(), "Node C started with OpenTelemetry tracing");
    }

private:
    void messageCallback(const std_msgs::msg::String::SharedPtr msg) {
        // Parse trace context from message
        std::string message_data = msg->data;
        std::string trace_id_hex, span_id_hex, original_message;
        
        // Extract trace information from message
        size_t trace_pos = message_data.find("TRACE:");
        size_t span_pos = message_data.find(":SPAN:");
        size_t msg_pos = message_data.find(":Processed by B:");
        
        if (trace_pos != std::string::npos && span_pos != std::string::npos && msg_pos != std::string::npos) {
            trace_id_hex = message_data.substr(trace_pos + 6, span_pos - trace_pos - 6);
            span_id_hex = message_data.substr(span_pos + 6, msg_pos - span_pos - 6);
            original_message = message_data.substr(msg_pos + 1);
        } else {
            // Fallback if trace info not found
            trace_id_hex = "00000000000000000000000000000000";
            span_id_hex = "0000000000000000";
            original_message = message_data;
        }
        
        // Create child span for message consumption (simplified approach)
        auto span = tracer_->StartSpan("consume_message", {
            {"node.name", "node_c"},
            {"topic.input", "topic_b_to_c"},
            {"message.type", "std_msgs/String"},
            {"operation.type", "consume"},
            {"parent.trace_id", trace_id_hex},
            {"parent.span_id", span_id_hex}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            RCLCPP_INFO(this->get_logger(), "Received: %s", original_message.c_str());
            RCLCPP_INFO(this->get_logger(), "Parent Trace ID: %s", trace_id_hex.c_str());
            RCLCPP_INFO(this->get_logger(), "Parent Span ID: %s", span_id_hex.c_str());
            
            // Add span attributes
            span->SetAttribute("input.message", original_message);
            span->SetAttribute("message.count", static_cast<int64_t>(message_count_));
            span->AddEvent("message_received");
            
            // Simulate processing time
            std::this_thread::sleep_for(std::chrono::milliseconds(150));
            
            // Process message (final processing)
            std::string result = "Final result from C: " + original_message;
            std::string current_span_id = "span_c_" + std::to_string(message_count_);
            
            span->SetAttribute("processing.result", result);
            span->SetAttribute("current.trace_id", trace_id_hex);
            span->SetAttribute("current.span_id", current_span_id);
            span->AddEvent("message_processed");
            
            // Simulate more processing time
            std::this_thread::sleep_for(std::chrono::milliseconds(120));
            
            // Simulate some final action
            span->AddEvent("final_action_completed");
            
            // Set success status
            span->SetStatus(trace_api::StatusCode::kOk, "Message consumed and processed successfully");
            
            RCLCPP_INFO(this->get_logger(), "Processed: %s", result.c_str());
            RCLCPP_INFO(this->get_logger(), "Current Trace ID: %s", trace_id_hex.c_str());
            RCLCPP_INFO(this->get_logger(), "Current Span ID: %s", current_span_id.c_str());
            // Note: ToHex() method may not be available in this version
            // RCLCPP_INFO(this->get_logger(), "Trace ID: %s", span->GetContext().trace_id().ToHex().c_str());
            // RCLCPP_INFO(this->get_logger(), "Span ID: %s", span->GetContext().span_id().ToHex().c_str());
            
            message_count_++;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            RCLCPP_ERROR(this->get_logger(), "Error processing message: %s", e.what());
        }
        
        span->End();
    }

    rclcpp::Subscription<std_msgs::msg::String>::SharedPtr subscription_;
    nostd::shared_ptr<trace_api::Tracer> tracer_;
    size_t message_count_ = 0;
};

int main(int argc, char* argv[]) {
    rclcpp::init(argc, argv);
    
    auto node = std::make_shared<NodeC>();
    
    RCLCPP_INFO(node->get_logger(), "Starting Node C with OpenTelemetry distributed tracing...");
    
    rclcpp::spin(node);
    rclcpp::shutdown();
    
    return 0;
}
